#include <opencv2/opencv.hpp>

#include "light_detect.h"
#include "param.h"

void ArmorDetect::BGR2Gray(cv::Mat &src, cv::Mat &gray) {
  cv::Mat channels[3];
  cv::split(src, channels);
  if (this->enemy == param::RED) {
    gray = channels[2] - channels[0];
  } else {
    gray = channels[0] - channels[2];
  }
}

void ArmorDetect::imageEnhance(cv::Mat &gray, cv::Mat &enhance) {
  // 二值化
  cv::Mat binary_image;
  cv::threshold(gray, binary_image, param::binary_thres, 255,
                cv::THRESH_BINARY);

  // 膨胀
  cv::Mat tmp_image;
  static cv::Mat element;
  element = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(3, 3));
  cv::dilate(binary_image, tmp_image, element);
  // 腐蚀
  element = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(3, 3));
  cv::erode(tmp_image, enhance, element);
  // 涂黑上侧区域，减少干扰
  cv::Rect roi(0, 0, enhance.cols, enhance.rows / 4);
  enhance(roi) = cv::Scalar(0);
  // cv::imshow("gray_img", enhance);
}

void ArmorDetect::findLightBar(cv::Mat &enhance,
                               std::vector<cv::RotatedRect> &light_bars) {
  std::vector<std::vector<cv::Point>> contours;
  cv::findContours(enhance, contours, cv::RETR_EXTERNAL,
                   cv::CHAIN_APPROX_SIMPLE);
  for (const auto &contour : contours) {
    if (contour.size() < 5) // 拟合椭圆至少需要 5 个点
      continue;

    cv::RotatedRect box = cv::fitEllipse(contour);
    double aspect_ratio = box.size.height * 1.0 / box.size.width;
    if (aspect_ratio < 1.0)
      aspect_ratio = 1.0 / aspect_ratio;

    if (aspect_ratio > param::max_light_ratio ||
        aspect_ratio < param::min_light_ratio)
      continue; // 长宽比筛选

    if (box.size.area() < param::min_light_area ||
        box.size.area() > param::max_light_area)
      continue; // 面积筛选

    // 1. 获取矩形的四个角点
    cv::Point2f vertices[4];
    box.points(vertices);

    // 2. 依次连接四个角点绘制四条边
    for (int i = 0; i < 4; i++) {
      // 使用 cv::line 函数连接点 i 和点 (i+1)%4
      cv::line(this->img, vertices[i], vertices[(i + 1) % 4],
               cv::Scalar(0, 255, 255), // 颜色：黄色 (BGR: 蓝0, 绿255, 红255)
               2);                      // 线条粗细
    }

    // std::string text = std::to_string(box.size.area()).substr(0, 5);
    // cv::putText(this->img, text, vertices[0] + cv::Point2f(-10, -20),
    //             cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(255, 255, 255), 1);
    light_bars.push_back(box);
  }
}

void ArmorDetect::matchLightBar(std::vector<cv::RotatedRect> &light_bars,
                                std::vector<Armor> &armors) {
  for (size_t i = 0; i < light_bars.size(); ++i) {
    for (size_t j = i + 1; j < light_bars.size(); ++j) {
      const auto &bar1 = light_bars[i];
      const auto &bar2 = light_bars[j];
      // 计算两个灯条的中心点距离
      double center_distance = cv::norm(bar1.center - bar2.center);
      // 计算两个灯条的平均高度
      double avg_height = (bar1.size.height + bar2.size.height) / 2.0;
      double height_diff = std::max(bar1.size.height, bar2.size.height) * 1.0 /
                           std::min(bar1.size.height, bar2.size.height);
      // 高度差筛选
      if (height_diff > param::max_armor_height_diff)
        continue;

      // 计算两个灯条的角度差
      double angle_diff = std::abs(bar1.angle - bar2.angle);
      if (angle_diff > param::max_angle_diff)
        continue; // 角度差筛选

      // 计算长宽比
      double length_ratio = center_distance / avg_height;
      if (length_ratio > param::max_armor_ratio ||
          length_ratio < param::min_armor_ratio)
        continue; // 长宽比筛选

      double area = center_distance * avg_height;
      if (area > param::max_area || area < param::min_area ||
          isnan(area)) // 注意判断 nan
        continue;      // 面积筛选

      // 获取装甲板图像
      // cv::Mat armor_img;
      // this->getArmorImg(this->src, {bar1, bar2}, armor_img);
      // auto [type, confidence_level] = logo_recognize.recognize(armor_img);
      // if (type == "Unknown")
      //   continue; // 兵种识别筛选

      // armors.push_back(Armor(bar1, bar2, type, confidence_level));
      armors.push_back(Armor(bar1, bar2));
      // std::cerr << area << std::endl;
    }
  }
}

void ArmorDetect::getArmorImg(const cv::Mat &src, const Armor &armor,
                              cv::Mat &armor_img) {
  cv::Point2f vertices[4];
  cv::RotatedRect armor_rect;
  const auto &[bar1, bar2, type, confidence_level] = armor;
  armor_rect.center = (bar1.center + bar2.center) / 2.0;
  double center_distance =
      cv::norm(bar1.center - bar2.center) * param::width_zoom_factor;

  // 高度缩放，得到整个装甲板
  double avg_height =
      (bar1.size.height + bar2.size.height) / 2.0 * param::length_zoom_factor;
  armor_rect.size = cv::Size2f(center_distance, avg_height);
  armor_rect.angle = (bar1.angle + bar2.angle) / 2.0;
  armor_rect.points(vertices);
  cv::Point2f target_vertices[4] = {
      cv::Point2f(0, param::target_height - 1),
      cv::Point2f(param::target_width - 1, param::target_height - 1),
      cv::Point2f(param::target_width - 1, 0), cv::Point2f(0, 0)};

  // 透视变换
  cv::Mat M = cv::getPerspectiveTransform(vertices, target_vertices);
  cv::warpPerspective(src, armor_img, M,
                      cv::Size(param::target_width, param::target_height));
}

void ArmorDetect::setEnemy(bool enemy) { this->enemy = enemy; }

void ArmorDetect::detect(cv::Mat &src, cv::Mat &dst) {
  cv::Mat gray, enhance;
  this->img = src;
  this->src = src.clone();
  this->BGR2Gray(src, gray);
  this->imageEnhance(gray, enhance);
  // cv::imshow("gray", enhance);

  std::vector<cv::RotatedRect> light_bars;
  this->findLightBar(enhance, light_bars);
  // cv::imshow("light", this->img);
  // cv::waitKey(0);

  std::vector<Armor> armors;
  this->matchLightBar(light_bars, armors);

  if (armors.size() > 0)
    this->getArmorImg(this->src, armors.front(), dst);
}